import numpy as np
import cv2
from scipy.optimize import linear_sum_assignment

class Track:
    def __init__(self, bbox, track_id, cls):
        self.bbox = bbox  # [x1,y1,x2,y2]
        self.id = track_id
        self.cls = cls
        self.missed = 0   # 丢失帧计数
        self.posivec = None  # 行驶方向向量
        self.center = ((bbox[0]+bbox[2])/2, (bbox[1]+bbox[3])/2)
        self.is_captured = False

    def update_center(self):
        self.center = ((self.bbox[0]+self.bbox[2])/2, (self.bbox[1]+self.bbox[3])/2)

def iou(b1, b2):
    x1, y1 = max(b1[0], b2[0]), max(b1[1], b2[1])
    x2, y2 = min(b1[2], b2[2]), min(b1[3], b2[3])
    inter = max(0, x2-x1) * max(0, y2-y1)
    area1 = (b1[2]-b1[0]) * (b1[3]-b1[1])
    area2 = (b2[2]-b2[0]) * (b2[3]-b2[1])
    union = area1 + area2 - inter
    return inter/union if union > 0 else 0

class ByteTrackLite:
    def __init__(self, iou_thresh=0.5, max_age=20, conf_thresh=0.55, roi_list=None):
        self.iou_thresh = iou_thresh
        self.max_age = max_age
        self.conf_thresh = conf_thresh
        self.tracks = {}  # key=track_id, value=Track对象
        self.next_id = 1
        self.roi_np = np.array(roi_list, dtype=np.int32) if roi_list is not None else None

    def _filter_detections_in_roi(self, detections):
        """
        检测框中心点是否在 ROI 内，返回布尔索引
        detections: np.array [N,6] → [x1,y1,x2,y2,score,cls]
        """
        if self.roi_np is None or len(detections) == 0:
            return np.ones(len(detections), dtype=bool)

        centers = np.column_stack(((detections[:,0]+detections[:,2])/2,
                                   (detections[:,1]+detections[:,3])/2))
        # 用 cv2.pointPolygonTest 批量判断
        mask = np.array([cv2.pointPolygonTest(self.roi_np, (x,y), False) >= 0
                         for x, y in centers])
        return mask

    def update(self, detections):
        """
        detections: np.array [N,6] → [x1,y1,x2,y2,score,cls]
        """
        if detections is None or len(detections) == 0:
            # 所有轨迹 missed++
            for t in self.tracks.values():
                t.missed += 1

            # 删除丢失过久轨迹
            removed_ids = [tid for tid, t in self.tracks.items() if t.missed > self.max_age]
            for tid in removed_ids:
                del self.tracks[tid]

            return list(self.tracks.values())
        # 只保留高置信度
        detections = detections[detections[:,4] > self.conf_thresh]

        # 筛选 ROI 内的检测框
        mask = self._filter_detections_in_roi(detections)
        detections = detections[mask]
        det_count = len(detections)
        if det_count == 0:
            return list(self.tracks.values())

        # 没有轨迹时直接生成
        if len(self.tracks) == 0:
            for det in detections:
                new_id = self.next_id
                self.next_id += 1
                self.tracks[new_id] = Track(det[:4], new_id, int(det[5]))
            return list(self.tracks.values())

        # 构建 IoU 矩阵
        track_ids = list(self.tracks.keys())
        iou_matrix = np.zeros((len(track_ids), det_count), dtype=np.float32)
        for t_idx, tid in enumerate(track_ids):
            t_bbox = self.tracks[tid].bbox
            for d_idx, det in enumerate(detections):
                iou_matrix[t_idx, d_idx] = iou(t_bbox, det[:4])

        # 匈牙利匹配
        row_ind, col_ind = linear_sum_assignment(-iou_matrix)
        unmatched_tracks = [tid for i, tid in enumerate(track_ids) if i not in row_ind]
        unmatched_dets = list(range(det_count))

        for r, c in zip(row_ind, col_ind):
            tid = track_ids[r]
            if iou_matrix[r, c] >= self.iou_thresh:
                prev_center = self.tracks[tid].center
                self.tracks[tid].bbox = detections[c,:4]
                self.tracks[tid].cls = int(detections[c,5])
                self.tracks[tid].missed = 0
                # 更新中心点
                self.tracks[tid].update_center()
                # 更新行驶方向向量
                self.tracks[tid].posivec = (self.tracks[tid].center[0]-prev_center[0],
                                            self.tracks[tid].center[1]-prev_center[1])
                unmatched_dets.remove(c)
            else:
                unmatched_tracks.append(tid)

        # 未匹配轨迹 missed++
        for tid in unmatched_tracks:
            self.tracks[tid].missed += 1

        # 删除丢失过久轨迹
        removed_ids = [tid for tid, t in self.tracks.items() if t.missed > self.max_age]
        for tid in removed_ids:
            del self.tracks[tid]

        # 新增轨迹
        for d_idx in unmatched_dets:
            det = detections[d_idx]
            new_id = self.next_id
            self.next_id += 1
            self.tracks[new_id] = Track(det[:4], new_id, int(det[5]))

        return list(self.tracks.values())

    # 根据 ID 访问车辆
    def get_track_by_id(self, track_id):
        return self.tracks.get(track_id, None)
